package consulTest.config;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.annotation.JSONField;
import lombok.Builder;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.util.ContentCachingRequestWrapper;
import org.springframework.web.util.ContentCachingResponseWrapper;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.function.BiConsumer;
import java.util.stream.Collectors;

@Slf4j
public class HTTPLogUtil {

    private static final Set<String> JSON_TYPE_SET = new HashSet<>(Arrays.asList("application/json;charset=UTF-8", "application/json"));

    public static void log(ContentCachingRequestWrapper request,
                           ContentCachingResponseWrapper response,
                           BiConsumer<ContentCachingRequestWrapper, ContentCachingResponseWrapper> biConsumer) {
        long start = System.nanoTime();
        biConsumer.accept(request, response);
        long end = System.nanoTime();

        int timeTaken = (int) ((end - start) / 1e6d);
        LogData logData = log(request, response, timeTaken,
                request.getContentAsByteArray(), response.getContentAsByteArray());
        log.info(JSON.toJSONString(logData));
    }

    private static LogData log(HttpServletRequest request,
                               HttpServletResponse response,
                               int timeTaken,
                               byte[] requestBody,
                               byte[] responseBody) {
        LogData.LogDataBuilder builder = LogData.builder()
                .path(request.getRequestURI())
                .method(request.getMethod())
                .remoteAddress(request.getRemoteAddr())
                .requestHeader(getRequestHeaders(request))
                .status(response.getStatus())
                .timeTaken(timeTaken);

        if (JSON_TYPE_SET.contains(request.getContentType())) {
            builder.requestBody(byte2String(requestBody));
        }
        if (JSON_TYPE_SET.contains(response.getContentType())) {
            builder.responseBody(byte2String(responseBody));
        }

        return builder.build();
    }

    private static Map<String, Object> getRequestHeaders(HttpServletRequest request) {
        Map<String, Object> headers = new HashMap<>();

        Enumeration<String> names = request.getHeaderNames();
        while (names.hasMoreElements()) {
            String name = names.nextElement();
            headers.put(name, getHeaderValue(Collections.list(request.getHeaders(name))));
        }
        return headers;
    }

    private static Map<String, Object> getResponseHeaders(HttpServletResponse response) {
        return response.getHeaderNames().stream()
                .collect(Collectors.toMap(a -> a, a -> getHeaderValue(response.getHeaders(a)), (o, n) -> o));
    }

    private static Object getHeaderValue(Collection<String> values) {
        if (Objects.isNull(values) || values.isEmpty()) {
            return "";
        }
        if (values.size() == 1) {
            return values.iterator().next();
        }
        return values;
    }

    private static String byte2String(byte[] bytes) {
        if (Objects.isNull(bytes) || bytes.length == 0) {
            return null;
        }
        return new String(bytes, StandardCharsets.UTF_8);
    }

    @Data
    @Builder
    public static class LogData {
        private String path;
        private String method;
        private String remoteAddress;

        private Object requestHeader;
        @JSONField(jsonDirect = true)
        private String requestBody;

        private Object responseHeader;
        @JSONField(jsonDirect = true)
        private String responseBody;

        private Integer status;
        private Integer timeTaken;
    }
}
